import os
import sys
import zipfile
from collections import defaultdict

def scan_zip_file(zip_path):
    """扫描ZIP文件，统计各类型文件的数量和占用空间"""
    ext_stats = defaultdict(lambda: {'count': 0, 'size': 0})
    total_count = 0
    total_size = 0

    try:
        with zipfile.ZipFile(zip_path, 'r') as zipf:
            for info in zipf.infolist():
                # 只统计文件（跳过目录）
                if not info.is_dir():
                    file_name = info.filename
                    file_size = info.file_size  # 使用未压缩大小（即原始大小）
                    
                    # 获取扩展名
                    _, ext = os.path.splitext(file_name)
                    if not ext:
                        ext = 'no_extension'
                    else:
                        ext = ext.lower()  # 统一转小写，避免 .TXT 和 .txt 被分开

                    # 更新该扩展名的统计
                    ext_stats[ext]['count'] += 1
                    ext_stats[ext]['size'] += file_size

                    # 更新总计
                    total_count += 1
                    total_size += file_size

    except zipfile.BadZipFile:
        print(f"错误：'{zip_path}' 不是一个有效的 ZIP 文件。")
        sys.exit(1)
    except FileNotFoundError:
        print(f"错误：找不到文件 '{zip_path}'。")
        sys.exit(1)
    except Exception as e:
        print(f"读取 ZIP 文件时发生未知错误: {e}")
        sys.exit(1)

    return dict(ext_stats), total_count, total_size


def format_bytes(size):
    """格式化字节大小为可读格式"""
    for unit in ['B', 'KB', 'MB', 'GB', 'TB']:
        if size < 1024:
            return f"{size:.2f} {unit}"
        size /= 1024
    return f"{size:.2f} PB"


if __name__ == "__main__":
    if len(sys.argv) != 2:
        print("使用方法: python zip_filetype.py <zip_file_path>")
        sys.exit(1)

    zip_file_path = sys.argv[1]

    stats, total_count, total_size = scan_zip_file(zip_file_path)

    print("ZIP 文件内文件类型统计：")
    for ext, info in sorted(stats.items()):
        count = info['count']
        size = format_bytes(info['size'])
        print(f"{ext}: 数量={count}, 占用空间={size}")
    
    print(f"\n总文件数量: {total_count}, 总占用空间（未压缩）: {format_bytes(total_size)}")